package main;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileReader;
import java.io.FileWriter;
import java.util.AbstractMap.SimpleEntry;
import java.util.ArrayList;

import translationModel.LanguageModel;
import translationModel.Phrase;
import translationModel.PhraseDecoder;
import translationModel.Sentence;

public class Main
{
	public static void main(String[] args) throws Exception 
	{
		InputArgs input = parseArgs(args);
		if (input == null)
			return;
		
		LanguageModel lm = ExtractLanguageModel(input.lmFile);
		ArrayList<Phrase> lexicon = ExtractPhraseLexicon(input.tmFile);
		
		System.out.println("Translation Started!");
		BufferedWriter outputFileWriter = new BufferedWriter(new FileWriter(input.outputFile));
		if (!input.forcedDecoding)
		{
			ArrayList<Sentence> inputSentences = ExtractInputSentences(input.inputFile);
			// Find the translation for each sentence
			for (Sentence sent : inputSentences)
			{
				PhraseDecoder decoder = new PhraseDecoder(sent, lexicon, lm, input.beamSize, input.distortionLimit);
				String output = decoder.RunAlgorithm();
				outputFileWriter.write(output + "\n");
			}
		}
		else
		{
			ArrayList<SimpleEntry<Sentence, Sentence>> inputSentences = ExtractForcedDecodingInputSentences(input.inputFile);
			// Find the forced decoding translation
			for (SimpleEntry<Sentence, Sentence> entry : inputSentences)
			{
				PhraseDecoder decoder = new PhraseDecoder(entry.getKey(), lexicon, lm, entry.getValue());
				String output = decoder.RunAlgorithm();
				outputFileWriter.write(output + "\n");
			}
		}
		outputFileWriter.close();
		System.out.println("Translation Done!");
	}

	private static ArrayList<Sentence> ExtractInputSentences(String inputFile) throws Exception
	{
		ArrayList<Sentence> sentences = new ArrayList<Sentence>();
		BufferedReader fileReader = new BufferedReader(new FileReader(inputFile));
		String line = fileReader.readLine();
		while(line != null)
		{
			sentences.add(new Sentence(line));
			line = fileReader.readLine();
		}
		fileReader.close();
		return sentences;
	}
	
	private static ArrayList<SimpleEntry<Sentence, Sentence>> ExtractForcedDecodingInputSentences(String inputFile) throws Exception
	{
		ArrayList<SimpleEntry<Sentence, Sentence>> sentences = new ArrayList<SimpleEntry<Sentence, Sentence>>();
		BufferedReader fileReader = new BufferedReader(new FileReader(inputFile));
		String f = fileReader.readLine();
		String e;
		while (f != null)
		{
			if (f.length() > 0)
			{
				e = fileReader.readLine();
				sentences.add(new SimpleEntry<Sentence, Sentence>(new Sentence(f), new Sentence(e)));
			}
			f = fileReader.readLine();
		}
		fileReader.close();
		return sentences;
	}

	private static ArrayList<Phrase> ExtractPhraseLexicon(String tmFile) throws Exception
	{
		ArrayList<Phrase> lexicon = new ArrayList<Phrase>();
		BufferedReader fileReader = new BufferedReader(new FileReader(tmFile));
		String line = fileReader.readLine();
		while(line != null)
		{
			String [] parts = line.split(" \\|\\|\\| ");
			lexicon.add(new Phrase(new Sentence(parts[0]), new Sentence(parts[1]), Double.parseDouble(parts[2])));
			line = fileReader.readLine();
		}
		fileReader.close();
		return lexicon;
	}

	private static LanguageModel ExtractLanguageModel(String lmFile) throws Exception 
	{
		LanguageModel lm = new LanguageModel();
		BufferedReader fileReader = new BufferedReader(new FileReader(lmFile));
		String line = fileReader.readLine();
		while(line != null)
		{
			if ((line.length() > 0) && (!line.startsWith("\\")) && (!line.startsWith("ngram")))
			{
				String[] lineWords = line.split("\t");
				String ngram = lineWords[1];
				double prob = Double.parseDouble(lineWords[0]);
				double backOffProb = (lineWords.length == 3) ? Double.parseDouble(lineWords[2]) : prob;
				lm.AddNGram(ngram, prob, backOffProb);
			}
			line = fileReader.readLine();
		}
		fileReader.close();
		return lm;
	}

	private static InputArgs parseArgs(String[] args)
	{
		InputArgs input = new InputArgs();
		for (int i = 0; i < args.length; i++)
		{
			if (args[i].equals("-lm"))
			{
				if (i == args.length - 1)
				{
					System.out.println("Please specify the lm file after the -lm");
					return null;
				}
				
				File file = new File(args[i + 1]);
				if ((!file.exists()) || (!file.canRead()))
				{
					System.out.println("Cannot read from the lm file");
					return null;
				}
				
				input.lmFile = args[i + 1];
				i++;
			}
			
			else if (args[i].equals("-tm"))
			{
				if (i == args.length - 1)
				{
					System.out.println("Please specify the tm file after the -tm");
					return null;
				}
				
				File file = new File(args[i + 1]);
				if ((!file.exists()) || (!file.canRead()))
				{
					System.out.println("Cannot read from the tm file");
					return null;
				}
				
				input.tmFile = args[i + 1];
				i++;
			}
			
			else if (args[i].equals("-i"))
			{
				if (i == args.length - 1)
				{
					System.out.println("Please specify the input file to translate after the -i");
					return null;
				}
				
				File file = new File(args[i + 1]);
				if ((!file.exists()) || (!file.canRead()))
				{
					System.out.println("Cannot read from the input file");
					return null;
				}
				
				input.inputFile = args[i + 1];
				i++;
			}
			
			else if (args[i].equals("-f"))
				input.forcedDecoding = true;
			
			else if (args[i].equals("-b"))
			{
				if (i == args.length - 1)
				{
					System.out.println("Please specify the beam size after the -b");
					return null;
				}
				
				int number;
				try
				{
					number = Integer.parseInt(args[i + 1]);
				}
				catch (NumberFormatException e)
				{
					System.out.println("The beam size should be a positive number");
					return null;
				}
				
				if (number < 1)
				{
					System.out.println("The beam size should be a positive number");
					return null;
				}
				
				input.beamSize = number;
				i++;
			}
			
			else if (args[i].equals("-d"))
			{
				if (i == args.length - 1)
				{
					System.out.println("Please specify the distortion limit after the -d");
					return null;
				}
				
				int number;
				try
				{
					number = Integer.parseInt(args[i + 1]);
				}
				catch (NumberFormatException e)
				{
					System.out.println("The distortion limit should be a positive number");
					return null;
				}
				
				if (number < 1)
				{
					System.out.println("The distortion limit should be a positive number");
					return null;
				}
				
				input.distortionLimit = number;
				i++;
			}
			
			else if (args[i].equals("-o"))
			{
				if (i == args.length - 1)
				{
					System.out.println("Please specify the output file after the -o");
					return null;
				}
				
				File file = new File(args[i + 1]);
				if ((file.exists()) && (!file.canWrite()))
				{
					System.out.println("Cannot write to the output file");
					return null;
				}
				
				input.outputFile = args[i + 1];
				i++;
			}
		}
		
		if ((input.lmFile == null) || (input.tmFile == null) || (input.inputFile == null))
		{
			System.out.println("Missing arguments");
			System.out.println("Use the following form:");
			System.out.println("-lm lm_file -tm tm_file -i input_file [ -f ] [ -b beam_size ] [ -d distortion_limit ] [ -o output_file ]");
			return null;
		}
		
		// Defaults
		if (input.forcedDecoding == null)
			input.forcedDecoding = false;
		if (input.beamSize == null)
			input.beamSize = 100;
		if (input.distortionLimit == null)
			input.distortionLimit = -1;
		if (input.outputFile == null)
			input.outputFile = "output.txt";
		
		return input;
	}

	private static class InputArgs
	{
		public String lmFile;
		public String tmFile;
		public String inputFile;
		public Boolean forcedDecoding;
		public Integer beamSize;
		public Integer distortionLimit;
		public String outputFile;
	}
}
